import numpy as np  
from scipy.integrate import tplquad  
import matplotlib.pyplot as plt  
  
class MassPoint:  
    def __init__(self, m, x, y, z):  
        self.m = m  
        self.x = x  
        self.y = y  
        self.z = z  
  
class Ellipsoid:  
    def __init__(self, m, a, b, c):  
        self.m = m  
        self.a = a  
        self.b = b  
        self.c = c  
  
def ellipsoid_v(ellipsoid, masspoint, G=6.67430e-11):  
    M = ellipsoid.m  
    m = masspoint.m  
      
    k = -(3 * G * M * m) / (4 * np.pi)  
      
    # 定义被积函数  
    def integrand(varphi, theta, r):  
        a, b, c = ellipsoid.a, ellipsoid.b, ellipsoid.c  
        x0, y0, z0 = masspoint.x, masspoint.y, masspoint.z  
          
        R = np.sqrt((a * r * np.cos(theta) * np.sin(varphi) - x0)**2 +  
                    (b * r * np.sin(theta) * np.sin(varphi) - y0)**2 +  
                    (c * r * np.cos(varphi) - z0)**2)  
          
        # 当 R 小于阈值时，返回 0  
        if R < 1e-2:  
            return 0  
        else:  
            return (r**2 * np.sin(varphi)) / R  
  
    # 进行三重积分  
    res, _ = tplquad(integrand, 0, np.pi, lambda x: 0, lambda x: 2*np.pi, lambda x, y: 0, lambda x, y: 1)  
    return k * res  
  
# 参数设置  
G = 1  
M = 1  
m = 1  
  
a = 1  
b = 1  
c = 1  
  
n = 10  
x = np.linspace(0, 3, n)  
  
ellipsoid = Ellipsoid(M, a, b, c)  
  
# 初始化引力势数组  
V = np.zeros(n)  
  
# 循环遍历x值，计算引力势  
for i in range(n):  
    mass_point = MassPoint(m, 0, 0, x[i])  
    V[i] = ellipsoid_v(ellipsoid, mass_point, G)  
    print(f"i: {i}, V: {V[i]}")  
  
# 绘制引力势随距离的变化图  
plt.plot(x, V)  
plt.xlabel('Distance x')  
plt.ylabel('Gravitational Potential V')  
plt.title('Gravitational Potential vs. Distance for an Ellipsoid')  
plt.show()